#!/usr/bin/env python3
import argparse
from datetime import datetime, timedelta
from cryptography import x509
from cryptography.hazmat.backends import default_backend
import mysql.connector
import sys
import re

def get_args():
    parser = argparse.ArgumentParser(description="南墙waf证书更新工具")
    parser.add_argument('-i', '--id', required=True, help="更新数据库条件的ID")
    parser.add_argument('-n', '--name', required=True, help="证书路径 (e.g., name.crt)")
    parser.add_argument('-k', '--key', required=True, help="私钥路径 (e.g., private.key)")
    parser.add_argument('-s', '--connection-string', required=True,
                       dest='connection_string', help="数据库连接字符串 (格式: user:password@host:port)")
    return parser.parse_args()

def parse_connection_string(conn_str):
    """解析MySQL连接字符串"""
    try:
        # 格式: user:password@host:port
        user_part, host_part = conn_str.split('@')
        user, password = user_part.split(':')
        host, port = host_part.split(':')
        return {
            'user': user,
            'password': password,
            'host': host,
            'port': int(port)
        }
    except Exception as e:
        print(f"解析连接字符串失败: {e}", file=sys.stderr)
        sys.exit(1)

def get_file_contents(file_path):
    try:
        with open(file_path, 'r') as f:
            return f.read()
    except IOError as e:
        print(f"获取文件失败: {e}", file=sys.stderr)
        sys.exit(1)

def check_cert_ssl(crt_str):
    try:
        cert = x509.load_pem_x509_certificate(crt_str.encode(), default_backend())
        not_after = cert.not_valid_after_utc
        # 转换为北京时间 (UTC+8)
        dt = not_after + timedelta(hours=8)
        return dt.strftime('%Y-%m-%d %H:%M:%S')
    except Exception as e:
        print(f"检查证书失败: {e}", file=sys.stderr)
        sys.exit(1)

def get_dns_names(cert):
    dns_names = []
    try:
        ext = cert.extensions.get_extension_for_class(x509.SubjectAlternativeName)
        san = ext.value
        for dns in san.get_values_for_type(x509.DNSName):
            dns_names.append(dns)
    except x509.ExtensionNotFound:
        pass  # 没有SAN扩展
    return dns_names

def main():
    args = get_args()
    print(f"ID: {args.id}")
    print(f"Name: {args.name}") 
    print(f"Key: {args.key}")
    print(f"Connection String: {args.connection_string}")

    crt_content = get_file_contents(args.name)
    key_content = get_file_contents(args.key)
    
    not_after = check_cert_ssl(crt_content)
    print(f"证书有效期: {not_after}")

    try:
        cert = x509.load_pem_x509_certificate(crt_content.encode(), default_backend())
        dns_names = str(get_dns_names(cert)).replace('\'', '"')
        print(f"证书域名: {dns_names}")

        # 解析连接字符串并连接MySQL
        conn_params = parse_connection_string(args.connection_string)
        conn = mysql.connector.connect(
            host=conn_params['host'],
            port=conn_params['port'],
            user=conn_params['user'],
            password=conn_params['password'],
            database='uuwaf'
        )
        cursor = conn.cursor()

        # 更新数据库
        update_query = """
            UPDATE `uuwaf`.`waf_certs` 
            SET sni=%s, crt=%s, `key`=%s, expired_at=%s, updated_at=NOW() 
            WHERE id=%s
        """
        cursor.execute(update_query, (str(dns_names), crt_content, key_content, not_after, args.id))
        conn.commit()
        print("更新数据库成功")

    except Exception as e:
        print(f"操作失败: {e}", file=sys.stderr)
        sys.exit(1)
    finally:
        if 'conn' in locals() and conn.is_connected():
            conn.close()

if __name__ == "__main__":
    main()
